#-------------------------------------------------------------------------------
# Name:        table_to_array.py
# Purpose:     Converts variables within PyTables table into a 2D array,
#              very similar to an SQL crosstab query.
#
# Author:      slha
#
#-------------------------------------------------------------------------------
import tables as tb
import numpy as np
import datetime
from dateutil import rrule
from matplotlib.dates import num2date, date2num
import time
import logging
from logging import debug, info
logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.DEBUG)
#from datatools import datatools

def main():

    d1 = datetime.datetime(2000, 1, 1, 1, 0, 0)
    num = date2num(d1)
    d2 = num2d(num)
    debug(d2)


    h5dir       = r'C:\UnBackedUp_Data\h5'
    table_fname = r'%s\kentish_new.h5' % h5dir
    array_fname = r'%s\kentish_array_new.h5' % h5dir
    table_name   = 'met'
    array_group  = 'met'
    array_fmode = 'w'

    debug(table_fname)
    debug(array_fname)
    info("converting variables in table %s/%s to array %s/%s" %(table_fname, table_name, array_fname,array_group ))

    #
    # Specifies the time-extent and interval between rows of the array.
    #
    start    = datetime.datetime(2007, 1, 1, 0, 0, 0)
    end      = datetime.datetime(2011, 4, 14, 0, 0, 0)
    interval = 10
    rec_rule = rrule.rrule(dtstart=start, until=end, freq=rrule.MINUTELY, interval=interval)
    debug("variables will be extracted from %s to %s on %sm inteval" %(start, end, interval))

    filters     = tb.Filters(complevel=3, complib="blosc")
    array_file  = tb.openFile(array_fname, array_fmode, filters=filters)
    table_file  = tb.openFile(table_fname, 'a')
    table       = table_file.getNode('/', table_name)

    #
    # Ensure the tbid and timestamp are indexed
    # Which will make things much faster later
    #
    # Using a completely sorted index, as other types
    # were causing exceptions to be thrown when reading them
    #
    if table.cols.tbid.index==None:         table.cols.tbid.createCSIndex()
    if table.cols.utctimestamp.index==None: table.cols.utctimestamp.createCSIndex()

    #numtimes = table.cols.utctimestamp.index.readSorted()
    #start = num2date(numtimes[0])
    #end   = num2date(numtimes[-1])

    array_group = array_file.createGroup(array_file.root, array_group)

    idx         = table.cols.tbid.index
    turbines    = sorted(list(set(idx.readSorted())))
    var_names   = table.colnames
    debug("table_to_array(table, array_file, array_group, rec_rule) called")
    debug("processing turbines: %s" % turbines)
    debug("processing vars: %s"     % var_names)

    # Don't treat these like ordinary variables
    # In NetCDF terms they are the coordinate variables
    if 'tbid' in var_names:
        var_names.remove('tbid')
    if 'utctimestamp' in var_names:
        var_names.remove('utctimestamp')

    table_to_array(table, turbines, var_names, array_file, array_group, rec_rule=rec_rule)

    debug("table_to_array done, closing h5 files")
    array_file.close()
    table_file.close()



def table_to_array(table, turbines, var_names, array_file, array_group, rec_rule=None, missing='NA'):
    """Read from a PyTables table and write hdf5 arrays to a separate file for quicker access.
    Similar to an SQL crosstab query.

    Arguments:
        @table       -- PyTables table to extract data from
        @turbines    -- a list of TurbineIDs to be processed
        @var_names   -- list of variable namess to process
        @array_file  -- an open PyTables file to write the array to
        @array_group -- PyTables group to add the array to
        @rec_rule    -- a dateutil rrule object specifiying which dates to process
                        if None, the first two entries in table for turbines[0] are used
                        for frequency, and first and last entry for start and end"""

    start_time = time.time()
    #print array_file

    datetimes  = list(rec_rule)     # continous list of datetimes
    datestrs   = [d.strftime('%Y-%m-%d %H:%M:%S') for d in datetimes]
    numtimes   = date2num(datetimes)
    nturbines  = len(turbines)      # number of turbines
    nvars      = len(var_names)     # number of variables to process
    nslots     = len(datetimes)     # number of timestamps over the whole period
    nrows      = nslots

    #debug(' %d timestamps in table '         % ntimes)
    debug(' %d timestamps will be processed' % nslots)
    debug( 'processing %d vars, %d turbines, %d rows' %(nvars, nturbines, nrows))

    #
    # Manually create array of turbine IDs and timestamps (coordinate arrays)
    #
    a = array_file.createArray(array_group, 'tbid', turbines, 'Turbine IDs')
    a =  array_file.createArray(array_group, 'utctimestamp',numtimes, 'UTC Timestamp')
    a =  array_file.createArray(array_group, 'utctimestring',datestrs, 'UTC Timestamp string')

    for var in var_names:
        data_array = np.zeros((nrows, nturbines), 'f')

        for j in range(nturbines):
            start_time  = time.time()
            turbine     = turbines[j]
            #debug(turbine)
            print 'processing variables %s at turbine %s' %(var, turbine)

            #
            # Build a dictionary of this variable, keyed on a datetime object from utctimestamp.
            # Make sure timestamp is >0 as screwy timestamps mess things up
            #
            #data_dict = dict([  strftime(x['utctimestamp']), x[var]) for x in table.iterrows() if x['tbid'] =='%s' %turbine ])
            data_dict  = dict([(strptime(x['utctimestamp']), x[var] )for x in table.where("tbid =='%s' "%turbine)])

            #
            # Fill up the data array, masking any missing rows and masking missing fields within rows
            # Missing data might be because the timestamp is missing, in which case we will get
            # a KeyError, or the timestamp is there but the data is already encoded as missing.
            #
            for i in range(nslots):
                try:
                    date = datetimes[i]
                    val  = data_dict[date]
                    if val==missing:
                        data_array[i,j] = np.ma.masked
                    else:
                        data_array[i,j] = val

                except KeyError:
                    #entire row missing
                    data_array[i,j] =  np.ma.masked

                #except ValueError:
                #    data_array[i,j] =  np.ma.masked

        #
        # Now write the data
        #
        a = array_file.createArray(array_group, var, data_array, var)
        array_file.flush()
        print "done in %0.3f seconds" % (time.time()- start_time)


def strptime(timestr):
    """ Specific implementation to convert timeformat into
    a python datetime object"""

    #d = datetime.datetime.strptime(timestr,'%Y-%m-%d %H:%M')
    y = int(timestr[0:4])
    m = int(timestr[5:7])
    d = int(timestr[8:10])
    h = int(timestr[11:13])
    mi = int(timestr[14:16])
    #se = int(timestr[17:19])
    se = 0
    return datetime.datetime(y,m,d,h,mi,se)

def num2d(num):
    """ Converts a number to a datetime object but drops the seconds and
    microseconds to avoid ambiguities"""

    if num<1:
        return None

    d1 = num2date(num)
    d2 = datetime.datetime(d1.year, d1.month, d1.day, d1.hour, d1.minute, 0, 0)
    return d2


def create_tables():
    #
    # Create table definition
    #
    use_cols,var_names, np_types, h5_types = datatools.create_table_def(control_file)
    from table_def import TurbineData
    h5tab      = tb.openFile(h5table, 'w')
    filters    = tb.Filters(complevel=3, complib="blosc")
    group      = h5tab.createGroup('/', group_name, 'scada data')
    table      = h5tab.createTable(group, table_name, TurbineData, title='turbine data', expectedrows=499522)

    #
    # Add data to table from csv files
    #
    datatools.add_data_to_table(table, data_files, use_cols, var_names, np_types)
    h5tab.close()

if __name__ == '__main__':
    main()
